
import subprocess

import json
import threading
import io
import os
import sys
import json
import socket
import logging
import logging.handlers
import signal
import glob
from types import SimpleNamespace
from concurrent.futures import ThreadPoolExecutor, as_completed
import ctypes
import queue
import time 
from time import sleep
from services import tftp
from services import dhcp
from services import helpers
from services import http
# ----------------------
# 全局配置
# ----------------------

args = None
pxeserver = None
dhcp_server = None
# default settings
SETTINGS = {'NETBOOT_DIR':'netboot',
            'NETBOOT_FILE':'',
            'DHCP_SERVER_IP':'192.168.2.2',
            'DHCP_SERVER_PORT':67,
            'DHCP_SERVER_NIC':'',
            'DHCP_OFFER_BEGIN':'192.168.2.100',
            'DHCP_OFFER_END':'192.168.2.150',
            'DHCP_SUBNET':'255.255.255.0',
            'DHCP_DNS':'8.8.8.8',
            'DHCP_ROUTER':'192.168.2.1',
            'DHCP_BROADCAST':'',
            'DHCP_FILESERVER':'192.168.2.2',
            'DHCP_WHITELIST':False,
            'HTTP_PORT':80,
            'HTTP_SERVER_IP':'0.0.0.0',
            'TFTP_SERVER_IP':'0.0.0.0',
            'DHCP_MODE_PROXY':False,
            'LOG_LEVEL':'',
            'MODE_VERBOSE':''}


# ----------------------
# 启动服务器
# ----------------------


from dataclasses import dataclass
from typing import Callable

@dataclass
class ServiceProfile:
    """服务配置模板"""
    target: Callable      # 服务启动函数
    name: str             # 服务名称
    port: int             # 监控端口
    protocol: str = 'tcp' # 监控协议
    max_retries: int = 5  
    backoff_base: float = 1.5
    max_backoff: float = 30.0

class MultiServiceMonitor:
    def __init__(self):
        self.lock = threading.RLock()
        self.service_pool: dict[str, ServiceProfile] = {}
        self.active_threads: dict[str, threading.Thread] = {}
        self.retry_states: dict[str, dict] = {}
        self._shutdown = threading.Event()
        self.msg_queue = queue.Queue()

    # ------------------- 通用健康检查方法 -------------------
    @staticmethod
    def _check_port(host: str, port: int, proto: str) -> bool:
        """端口连通性检查"""
        sock_type = socket.SOCK_STREAM if proto == 'tcp' else socket.SOCK_DGRAM
        with socket.socket(socket.AF_INET, sock_type) as s:
            s.settimeout(2)
            try:
                if proto == 'tcp':
                    s.connect((host, port))
                else:  # UDP检测发送空包
                    s.sendto(b'', (host, port))
                return True
            except Exception:
                return False

    # ------------------- 服务管理核心逻辑 -------------------
    def add_service(self, profile: ServiceProfile):
        """注册新服务"""
        with self.lock:
            self.service_pool[profile.name] = profile
            self.retry_states[profile.name] = {
                'retries': 0,
                'next_attempt': 0,
                'last_error': None
            }
            self._start_service(profile.name)

    def _start_service(self, name: str):
        """启动服务线程"""
        profile = self.service_pool[name]
        
        def _service_wrapper(msg_queue):
            try:
                sys_logger.info(f"[{name}] 服务启动中...")
                profile.target(msg_queue)
            except OSError as e:
                if e.errno == 9:
                    sys_logger.warning(f"[{name}] 检测到安全关闭")
                else:
                    sys_logger.error(f"[{name}] 套接字错误: {str(e)}")
            except Exception as e:
                sys_logger.error(f"[{name}] 服务崩溃: {str(e)}", exc_info=True)
                self._handle_failure(name, e)
            finally:
                with self.lock:
                    if name in self.active_threads:
                        del self.active_threads[name]
        thread = threading.Thread(
            target=_service_wrapper,
            name=f"Service-{name}",
            daemon=True,
            args=(self.msg_queue,)

        )
        with self.lock:
            self.active_threads[name] = thread
        thread.start()

    def _handle_failure(self, name: str, error: Exception):
        """失败处理策略"""
        with self.lock:
            state = self.retry_states[name]
            state['retries'] += 1
            state['last_error'] = str(error)
            
            # 计算指数退避时间
            delay = min(
                self.service_pool[name].backoff_base ** state['retries'],
                self.service_pool[name].max_backoff
            )
            state['next_attempt'] = time.time() + delay
            sys_logger.warning(f"[{name}] {delay:.1f}秒后尝试第{state['retries']}次重启")

    def _needs_restart(self, name: str) -> bool:
        """判断是否满足重启条件"""
        profile = self.service_pool[name]
        state = self.retry_states[name]
        
        return all([
            state['retries'] < profile.max_retries,
            time.time() > state['next_attempt'],
            not self._check_port('127.0.0.1', profile.port, profile.protocol)
        ])

    def run_monitor(self):
        """监控主循环"""
        sys_logger.info("====== 启动多服务监控系统 ======")
        try:
            while not self._shutdown.is_set():
                try:
                    #进程间通信
                    msg = self.msg_queue.get(timeout=0.1)
                    if msg =="restart":
                        monitor.restart_all_services()
                except:
                    pass
                with self.lock:
                    # 检查所有服务的健康状态
                    for name in list(self.service_pool.keys()):
                        # 服务未运行且需要重启
                        if name not in self.active_threads and self._needs_restart(name):
                            sys_logger.info(f"[{name}] 触发自动重启")
                            self._start_service(name)
                            
                    # 强制关闭条件检查
                    if any(state['retries'] >= self.service_pool[name].max_retries 
                           for name, state in self.retry_states.items()):
                        sys_logger.critical("达到最大重试次数，系统关闭")
                        break
                sleep(1)
                
        except KeyboardInterrupt:
            sys_logger.info("接收到关闭信号")
            self._shutdown_system()
        finally:
            try:
                self._shutdown_system()
            except:
                pass

    def _shutdown_system(self):
        """系统关闭流程"""
        with self.lock:
            sys_logger.info("正在停止所有服务...")
            monitor._graceful_shutdown()
            monitor._cleanup_resources()
        raise KeyboardInterrupt

class EnhancedMonitor(MultiServiceMonitor):
    def __init__(self):
        super().__init__()
        self.restart_lock = threading.Lock()
        self.last_restart = 0
        self._register_signals()

    def _register_signals(self):
        signal.signal(signal.SIGHUP, self._graceful_shutdown)  # SSH断开
        signal.signal(signal.SIGINT, self._graceful_shutdown)  # Ctrl+C
        signal.signal(signal.SIGTERM, self._graceful_shutdown) # kill命令
        
    def restart_all_services(self):
        """无验证的全局重启操作"""
        if not self._acquire_restart_lock():
            raise RuntimeError("Another restart is in progress")
            
        try:
            sys_logger.info("===== 启动全局服务重启 =====")
            self._graceful_shutdown()
            self._cleanup_resources()
            print("正在重启服务")
            output = subprocess.check_output(['bash', 'scripts/gen_boot_ipxe.sh'], text=True)
            self._parallel_restart()
            self.last_restart = time.time()
            return True
        finally:
            self.restart_lock.release()

    def _acquire_restart_lock(self):
        """非阻塞锁获取"""
        return self.restart_lock.acquire(blocking=False)

    def _graceful_shutdown(self):
        """优雅停止所有服务"""
        with self.lock:
            sys_logger.info("停止运行中的服务...")
            for name in list(self.service_pool.keys()):
                if hasattr(self.service_pool[name].target, '__self__'):
                    service = self.service_pool[name].target.__self__
                    if hasattr(service, 'stop'):
                        service.stop()
            for name, thread in self.active_threads.copy().items():
                if thread.is_alive():
                    sys_logger.info(f"正在停止服务 [{name}]")
                    self._force_stop_thread(thread)
            self.active_threads.clear()

    def _force_stop_thread(self, thread):
        """强制终止线程"""
        try:
            from ctypes import pythonapi, py_object
            pythonapi.PyThreadState_SetAsyncExc(
                ctypes.c_long(thread.ident),
                ctypes.py_object(SystemExit)
            )
        except Exception as e:
            sys_logger.error(f"终止线程失败: {str(e)}")

    def _cleanup_resources(self):
        """清理系统资源"""
        sys_logger.info("清理网络连接和文件句柄...")
        # 此处添加具体的资源回收逻辑
        if hasattr(self, '_network_connections'):
            for conn in self._network_connections.values():
                conn.close()

    def _parallel_restart(self):
        """并行重启优化"""
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = {}
            for name in self.service_pool:
                future = executor.submit(self._safe_restart_service, name)
                futures[future] = name
                
            for future in as_completed(futures):
                name = futures[future]
                try:
                    future.result()
                except Exception as e:
                    sys_logger.error(f"服务 [{name}] 重启失败: {str(e)}")

    def _safe_restart_service(self, name):
        """带异常保护的服务重启"""
        try:
            self._start_service(name)
            sys_logger.info(f"服务 [{name}] 重启成功")
            return True
        except Exception as e:
            sys_logger.error(f"服务 [{name}] 启动异常: {str(e)}")
            raise


if __name__ == "__main__":
    print("####开始启动服务####")
    print("正在扫描iso目录")
    iso_pattern = os.path.join("./iso", '*.iso')
    if not glob.glob(iso_pattern):
        print("iso文件夹中未发现镜像文件，请将镜像文件上传至iso文件夹后重新运行服务。")
    output = subprocess.check_output(['bash', 'scripts/gen_boot_ipxe.sh'], text=True)
    try:
        config_file = io.open("config/config.json", 'r')
    except IOError:
        sys.exit('Failed to open {0}'.format(args.JSON_CONFIG))
    try:
        loaded_config = json.load(config_file)
        config_file.close()
    except ValueError:
        sys.exit('{0} does not contain valid JSON'.format(args.JSON_CONFIG))
    for setting in loaded_config:
        if type(loaded_config[setting]) is bytes:
            loaded_config[setting] = loaded_config[setting].encode('ascii')
    SETTINGS.update(loaded_config) # update settings with JSON config
    args = SimpleNamespace(**SETTINGS)
    sys_logger = logging.getLogger('M-PXE')
    log_file = "logs/pxe.log"
    handler = logging.handlers.RotatingFileHandler(log_file, maxBytes=10 * 1024 * 1024, backupCount=5)
    formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s %(message)s')
    handler.setFormatter(formatter)
    sys_logger.addHandler(handler)
    log_levels = {
        'DEBUG': logging.DEBUG,
        'INFO': logging.INFO,
        'WARNING': logging.WARNING,
        'ERROR': logging.ERROR,
        'CRITICAL': logging.CRITICAL
    }
    sys_logger.setLevel(log_levels.get(args.LOG_LEVEL.upper(), logging.INFO))
    monitor = EnhancedMonitor()
    #重启NFS
    print("正在启动nfs")
    nfs_file_path = '/etc/exports'
    if not os.path.exists(nfs_file_path):
        print("请先安装nfs服务：sudo yum install nfs-utils 或者sudo apt install nfs-kernel-server")
    else:
        try:
            with open(nfs_file_path, 'r', encoding='utf-8') as file:
                lines = file.readlines()
            with open(nfs_file_path, 'w', encoding='utf-8') as file:
                for line in lines:
                    if '###Misstar' not in line:
                        file.write(line)
                file.write("/tmp/pxenfsroot *(ro,no_subtree_check,crossmnt,fsid=0)  ###MisstarPXE\n")
        except:
            print("修改/etc/exports失败，请关闭应用保护")
        try:
            # 使用systemctl命令重启NFS服务 systemctl reload nfs && systemctl start nfs
            subprocess.run(['systemctl', 'reload','nfs-server'])
            subprocess.run(['systemctl','restart', 'nfs-server'], check=True)
        except subprocess.CalledProcessError as e:
            print("重启NFS服务时出错:", e)
    print("正在启动tftp")
    tftp_logger = helpers.get_child_logger(sys_logger, 'TFTP')
    tftp_server = tftp.TFTPD(
        log_level = args.LOG_LEVEL,
        logger = tftp_logger,
        netboot_directory = args.NETBOOT_DIR,
        port = args.TFTP_PORT,
        ip = args.TFTP_SERVER_IP)
    monitor.add_service(ServiceProfile(
        name="TFTP-Server",
        target=tftp_server.listen,
        port=69,
        protocol='udp',
        max_retries=3
    ))
    print("正在启动dhcp")
    dhcp_logger = helpers.get_child_logger(sys_logger, 'DHCP')
    dhcp_server = dhcp.DHCPD(
            nic=args.DHCP_SERVER_NIC,
            ip = args.DHCP_SERVER_IP,
            port = args.DHCP_SERVER_PORT,
            offer_from = args.DHCP_OFFER_BEGIN,
            offer_to = args.DHCP_OFFER_END,
            subnet_mask = args.DHCP_SUBNET,
            router = args.DHCP_ROUTER,
            dns_server = args.DHCP_DNS,
            broadcast = args.DHCP_BROADCAST,
            file_server = args.DHCP_FILESERVER,
            file_name = args.NETBOOT_FILE,
            use_ipxe = args.USE_IPXE,
            use_http = args.USE_HTTP,
            mode_proxy = args.DHCP_MODE_PROXY,
            log_level = args.LOG_LEVEL,
            whitelist = args.DHCP_WHITELIST,
            logger = dhcp_logger,
            saveleases = args.LEASES_FILE)
    monitor.add_service(ServiceProfile(
        name="DHCP-Server",
        target=dhcp_server.listen,
        port=67,
        protocol='udp',
        max_retries=3
    ))
    print("正在启动http")
    http_logger = helpers.get_child_logger(sys_logger, 'HTTP')
    httpserver=http.HTTPServer("0.0.0.0",80,log_level = args.LOG_LEVEL,logger = http_logger)
    monitor.add_service(ServiceProfile(
        name="HTTP-Server",
        target=httpserver.start,
        port=80,
        max_retries=5,
        backoff_base=2
    ))
    print("####服务启动完成####")
    print(f"管理后台：http://{args.DHCP_SERVER_IP}/")
    monitor.run_monitor()